#!/usr/bin/env python -W ignore

#
# Copyright (c) 2010-2011, Davide Cittaro
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the <organization> nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#


import dclib
import numpy as np

def main():
  # get command line options and print them
  options = dclib.IO.parseOptions(dclib.IO.prepareOptions())
  dclib.IO.debugParameters(options) 

  # get chromosome length. This is useful as we treat chromosome as fixed size arrays
  # we need to allocate space...
  chromLength = dclib.IO.getChromosomeSizes(options)
  chromNames = chromLength.keys()
  chromNames.sort()

  # open handlers for output files, if needed
  if not options.noprof:
    bdgh = dclib.IO.openBedGraph(options)
  if options.findPeaks:
    peaksh = dclib.IO.openPeaks(options)   

  # Here starts everything: a giant loop over chromosomes!
  # basically we scan the pipeline options and we repeat that for each chromosome
  for chrom in chromNames:
    theSignalA = np.array([])	# signal A, typically the ChIP
    theSignalB = np.array([])	# signal B, typicall the Input or the mock
    options.zerodata = True
    options.onesignal = False
    
    dclib.Helpers.Debug(1, dclib.Helpers.timestring(), "Processing chromosome", chrom)
    dclib.Helpers.Debug(1, dclib.Helpers.timestring(), "Reading", options.sigA)
    theSignalA = dclib.IO.inputWrapper(chrom, chromLength[chrom], options, 1)
    if len(theSignalA) == 0:
      dclib.Helpers.Debug(1, "There is no data for", chrom, "in signal A, skipping")
      continue
    if options.sigB:
      dclib.Helpers.Debug(1, dclib.Helpers.timestring(), "Reading", options.sigB)
      theSignalB = dclib.IO.inputWrapper(chrom, chromLength[chrom], options, 0)
      if len(theSignalB) == 0:
        dclib.Helpers.Debug(1, "There is no data for", chrom, "in signal B, skipping")
        continue

    # From this point on, every step may be saved into a npy object, to load in ipython
    # and debug what's wrong...
    if options.npysave:
      dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "RAW")
      if options.sigB and not options.onesignal:
        dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "RAW")

    # collect data for possible scaling
    if options.scale:
      dataPTP = np.ptp(theSignalA)

    # Begin to loop over pipeline steps. Note that pipeline is a string and we exploit
    # python string iterator.
    for analysisStep in options.pipeline:
      if analysisStep == 'W':
        # Wavelet denoising
        theSignalA = dclib.NumAnalysis.PPLWavelet(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLWavelet(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "wltd")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "wltd")
      if analysisStep == 'A':
        # Autocorrelation
        theSignalA = dclib.NumAnalysis.PPLAutocorrelate(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLAutocorrelate(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "ac")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "ac")
      if analysisStep == 'F':
        # FIR Smoothing
        theSignalA = dclib.NumAnalysis.PPLSmooth(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLSmooth(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "Smooth")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "Smooth")
      if analysisStep == 'G':
        # Gaussian Smoothing
        theSignalA = dclib.NumAnalysis.PPLGaussSmooth(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLGaussSmooth(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "GS")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "GS")
      elif analysisStep == 'E':
        # Signal equalization
        theSignalA = dclib.NumAnalysis.PPLequalize(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLequalize(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "EQ")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "EQ")
      elif analysisStep == 'N':
        # Signal normalization
        theSignalA = dclib.NumAnalysis.PPLnormalize(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLnormalize(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "Norm")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "Norm")
      elif analysisStep == 'S':
        # Subtract signal B from signal A 
        theSignalA = dclib.NumAnalysis.PPLsubtract(theSignalA, theSignalB, options)
        del(theSignalB)
        # delete options.sigB so that this can be a step in the middle of 
        # the pipeline. Now we are running with a single dataset
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "diff")
      elif analysisStep == 'X' or analysisStep == 'V':
        # Cross-correlate or convolve...
        theSignalA = dclib.NumAnalysis.PPLXCorr(theSignalA, theSignalB, options)
        del(theSignalB)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "xcorr")
      elif analysisStep == 'R':
        # Calculate data ratio, i.e. (A-B)/A 
        theSignalA = dclib.NumAnalysis.PPLRatio(theSignalA, theSignalB, options)
        del(theSignalB)
        # As above, reduce to a single dataset
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "ratio")
      elif analysisStep == 'M':
        # Multiply
        theSignalA = dclib.NumAnalysis.PPLMult(theSignalA, theSignalB, options)
        del(theSignalB)
        # As above, reduce to a single dataset
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "multip")
      elif analysisStep == 'L':
        # Calculate log2ratio between A and B. This is, by far, the most common step
        theSignalA = dclib.NumAnalysis.PPLL2R(theSignalA, theSignalB, options)
        del(theSignalB)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "l2r")
      elif analysisStep == 'T':
        # Perform baseline correction. This step needs some refinement...
        theSignalA = dclib.NumAnalysis.PPLThreshold(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLThreshold(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "bsl")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "bsl")
      elif analysisStep == 'Z':
        # Remove negative values. Used when we want a 1-tail p-value estimation for peaks
        # and also when we are interested only in signal A properties
        theSignalA = dclib.NumAnalysis.PPLZero(theSignalA, options)
        if not options.onesignal and options.sigB:
          theSignalB = dclib.NumAnalysis.PPLZero(theSignalB, options)
        if options.npysave:
          dclib.IO.npysave(theSignalA, options.name, chrom, "SignalA", "zero")
          if options.sigB and not options.onesignal:
            dclib.IO.npysave(theSignalB, options.name, chrom, "SignalB", "zero")

    if options.correlate and not options.onesignal:
      # Here we calculate correlation between signal A and signal B. This will
      # plot a chart and linear regression data.
      corr_results = dclib.NumAnalysis.PPLCorrelate(theSignalA, theSignalB, options)
      ccoef = corr_results[2][0]
      cpvalue = corr_results[2][1]
      if options.npysave:
        dclib.IO.npysave(corr_results[0], options.name, chrom, "xc1")
        dclib.IO.npysave(corr_results[1], options.name, chrom, "xc2")
      dclib.IO.saveCorrelation(corr_results, chrom, options)

    # collect data for possible scaling
    if options.scale:
      theSignalA = dclib.NumAnalysis.PPLScale(theSignalA, dataPTP)


    if options.savestats:
      # save basic statistics. This is quite useless :-)
      dclib.IO.saveStats(theSignalA, chrom, options)

    if not options.noprof:
      # Save Signal into a the profile... Actually only signal A is saved. @@FIXTHIS@@
#      if options.profileFormat == 'bdg':
      dclib.IO.writeBedGraph(bdgh, theSignalA, options, chrom=chrom)
#      elif options.profileFormat == 'bw':
#        dclib.IO.writeBigWig(theSignalA, options, chromLength, chrom=chrom)

    if options.findPeaks:
      # Save peaks into the .peaks file. Only peaks in signal A will be retained. @@FIXTHIS@@
      dclib.Helpers.Debug(1, dclib.Helpers.timestring(), "Finding peaks boundaries")
      peaks = dclib.NumAnalysis.findPeaks(theSignalA, options)
      if not len(peaks):
        dclib.Helpers.Debug(1, dclib.Helpers.timestring(), "No peaks for chromosome", chrom)
        continue
      dclib.Helpers.Debug(1, dclib.Helpers.timestring(), "Writing peak data in peaks file")
      dclib.IO.writePeaks(peaksh, theSignalA, peaks, options, chrom=chrom)
      if not options.onesignal:
        dclib.Helpers.Debug(1, "WARNING: we are not writing peaks for signal B")

    # Explicitly free some memory...
    del(theSignalA)
    if not options.onesignal:
      del(theSignalB)

  # At the end of the loop it's time to close file handlers...
  if not options.noprof:
    if options.profileFormat == 'bdg':  
      dclib.IO.closeBedGraph(bdgh)   
    elif options.profileFormat == 'bw':
      dclib.IO.bdg2bw(bdgh, chromLength, options)
  if options.findPeaks:
    dclib.IO.closePeaks(peaksh)
  
if __name__ == "__main__": main()
